from __future__ import print_function
from numpy import *
import numpy as np
import sys
import scipy.io
import scipy

from scipy.optimize import fmin_powell

#####################################################################
## This gets the tree last parameters from the quaternion and maps it
## into values that make sense. Either just to the unit sphere, or
## eliminating the symmetries...
def fix_quaternion_parameters2(myx):
  xb,xc,xd = myx
  xnormsq = xb*xb+xc*xc+xd*xd

  if xnormsq < 1:
    ## If inside the unit sphere, these are the components
    ## themselves, and we just have to calculate a to normalize
    ## the quaternion.
    b,c,d = xb,xc,xd
    a = np.sqrt(1-xnormsq)
  else:
    ## Just to work gracefully if we have invalid inputs, we
    ## reflect the vectors outside the unit sphere to the other
    ## side, and use the inverse norm and negative values of a.
    b,c,d = -xb/xnormsq,-xc/xnormsq,-xd/xnormsq
    a = -np.sqrt(1- 1.0/xnormsq )
    ## This should not be used, it defeats the whole concept that
    ## small (b,c,d) vector norms have small rotation angles. It's
    ## really just to let us work near the borders of the
    ## sphere. Any optimization algorithm should work with
    ## initalizations just inside the spere, and avoid wandering
    ## outside of it.

  assert a >= -1
  assert a <= 1
  
  return a,b,c,d
##
#####################################################################

#####################################################################
## Produces a rotation matrix from the 3 last components of a
## quaternion.
def quaternion_to_matrix(myx):
  '''Converts from a quaternion representation (last 3 values) to rotation matrix.'''

  a,b,c,d = fix_quaternion_parameters2(myx)
  
  ## Notice we return a transpose matrix, because we work with line-vectors
  return array([ [(a*a+b*b-c*c-d*d), (2*b*c-2*a*d),   (2*b*d+2*a*c)   ],
             [(2*b*c+2*a*d),   (a*a-b*b+c*c-d*d), (2*c*d-2*a*b)   ],
             [(2*b*d-2*a*c),   (2*c*d+2*a*b),   (a*a-b*b-c*c+d*d)] ] ).T \
             / (a*a+b*b+c*c+d*d)


##
#####################################################################

def func(q, M):
    R = quaternion_to_matrix(q)
    return np.sum(np.abs(R-M))

def matrix_to_quaternion(Q):
    T = 1. + trace(Q)

    if T > 1e-8:
        W = sqrt(T/4.)
        X = (Q[1,2]-Q[2,1]) / (4. * W)
        Y = (Q[2,0]-Q[0,2]) / (4. * W)
        Z = (Q[0,1]-Q[1,0]) / (4. * W)

        # print(W*W+X*X+Y*Y+Z*Z)
        nf = (W*W+X*X+Y*Y+Z*Z)**-.5
        return array((X,Y,Z))*nf

    # else
    perm = array([[0,1,2], [1,2,0], [2,0,1]])
    p = perm[argmax(diag(Q))]

    u,v,w = p

    xx = 1 + Q[u,u] - Q[v,v] - Q[w,w]

    if xx <= 0:
        return zeros(3)

    r = np.sqrt(xx)
    
    q0 = (Q[w,v] - Q[v,w]) / (2 * r)
    qu = r / 2
    qv = (Q[u,v] + Q[v,u]) / (2 * r)
    qw = (Q[w,u] + Q[u,w]) / (2 * r)

    nf = (q0*q0+qu*qu+qv*qv+qw*qw)**-0.5


    out = zeros(3)
    out[p] = array([qu,qv,qw])*nf
    
    return out

if __name__ == '__main__':
  set_printoptions(precision=3)

  

  Rs = [array([[ 1, 0, 0],
               [ 0, 1, 0],
               [ 0, 0, 1] ]),
        array([[ 1, 0, 0],
               [ 0, 0, 1],
               [ 0,-1, 0] ]),
        array([[ 0, 0, 1],
               [ 0, 1, 0],
               [-1, 0, 0] ]),
        array([[ 0, 1, 0],
               [-1, 0, 0],
               [ 0, 0, 1] ]),
        ]

  rr = set([tuple(dot(dot(x,y),z).ravel()) for x in Rs  for y in Rs  for z in Rs])
  rr = list(rr)

  rr.sort()
  rr.reverse()

  rr = [reshape(array(r),(3,3)) for r in rr]

  def normalize(M):
      L = [dot(M,r) for r in rr]

  root = '/home/nlw/ciencia/DADOS/dists/YorkUrbanDB/'
  name = sys.argv[1]

  M = scipy.io.loadmat(root+name+'/'+name+'GroundTruthVP_Orthogonal_CamParams.mat')['vp_orthogonal']
  #print(name)
  #print(M)
  #print()
  angs = array( [(1-np.sum(matrix_to_quaternion(dot(M,r))**2)) for r in rr])
  zz = argmax(angs)
  print(name, *list(matrix_to_quaternion(dot(M,rr[zz]))))
